import numpy as np
import os
from PIL import Image
import mindspore as ms
from pathlib import Path
from scipy.io import loadmat


"""
3D Relightable Facial Expression (ICT-3DRFE)

Construct a dataset from
    Exploring the effect of illumination on automatic expression recognition using the ICT-3DRFE database
    https://www.sciencedirect.com/science/article/pii/S0262885612000169
"""
class Dataset(object):
    def __init__(self, name, root, batch_size=1):
        self.name = name
        self.root = Path(root)
        self.batch_size = batch_size
        self.image_extension = 'png'
        self.images_root = 'images'

    def rescale_image(self, image, method=0):
        return _rescale_image(image, method=0)

    def num_samples(self):
        return len(self._keys)

    def get_keys(self):
        path = self.root / self.images_root
        keys = [str(x.stem) for x in path.glob('*')]
        self._keys = keys

        print('Found {} files.'.format(len(keys)))

        if len(keys) == 0:
            raise RuntimeError('No images found in {}'.format(path))
        return keys

    def preprocess(self, image):
        return caffe_preprocess(image)

    def get_images(self, index, shape=None):
        return self._get_images(
            index,
            shape=None,
            subdir=self.images_root,
            extension=self.image_extension)

    def _get_images(self,
                    index,
                    shape=None,
                    subdir='images',
                    channels=3,
                    extension='png'):
        path = os.path.join(str(self.root / subdir), '/', index, '.', extension)

        image = Image.open(path).convert("RGB")

        return image

    def get_normals(self, index, shape=None):
        path = self.root / 'normals' / "{}.mat".format(
            index.decode("utf-8"))

        if path.exists():
            mat = loadmat(str(path))
            normals = mat['norms'].astype(np.float32)
            mask = mat['cert'][..., None].astype(np.float32)
            return normals, mask

        normals = np.zeros(shape).astype(np.float32)
        mask = np.zeros(list(shape[:2]) + [1]).astype(np.float32)
        return normals, mask


    def get_segmentation(self, index, shape=None):
        res = ms.ops.Zeros(shape)
        return res, res


class ICT3DFE(Dataset):
    def __init__(self, batch_size=1):
        super().__init__(
            name='ICT3DFE', root=Path('data/ict3drfe/'), batch_size=batch_size)
        self.image_extension = 'png'

    def get_normals(self, index, shape=None):
        normals, mask = super().get_normals(index, shape)
        return normals * np.array([-1, -1, 1]), mask


def caffe_preprocess(image):
    VGG_MEAN = np.array([102.9801, 115.9465, 122.7717])
    # RGB -> BGR
    image = ms.ops.reverse_sequence(image, [False, False, True])
    # Subtract VGG training mean across all channels
    image = image - VGG_MEAN.reshape([1, 1, 3])
    return image


def _rescale_image(image, stride_width=64, method=0):
    # make sure smallest size is 600 pixels wide & dimensions are (k * stride_width) + 1
    height = ms.ops.Shape(image)[0]
    width = ms.ops.Shape(image)[1]

    # Taken from 'szross'
    scale_up = 625. / ms.ops.minimum(height, width)
    scale_cap = 961. / ms.ops.maximum(height, width)
    scale_up = ms.ops.minimum(scale_up, scale_cap)
    new_height = stride_width * ms.ops.round(
        (height * scale_up) / stride_width) + 1
    new_width = stride_width * ms.ops.round((width * scale_up) / stride_width) + 1
    image = ms.dataset.vision.Resize(
        image, (new_height, new_width), method=method)
    return image



